Summary of Changes

Hello, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request addresses a bug in the layer specification logic for dense models within the Megatron framework. The changes ensure that the num_experts parameter is correctly interpreted, preventing unintended activation of Mixture-of-Experts (MoE) grouped GEMM operations when experts are not configured or explicitly set to zero. This fix improves the robustness and accuracy of model configuration, particularly for non-MoE architectures.

Highlights

• Bug Fix: Dense Model Layer Specification: Corrected the logic for determining the num_experts and moe_grouped_gemm parameters within the get_layer_spec function. Previously, num_experts defaulted to 0, which could incorrectly enable MoE grouping for dense models or misinterpret the absence of explicit expert configuration.
• Improved num_experts Handling: Changed the default value for num_experts from 0 to None when it's not explicitly provided in mg_config_dict. This allows for a clearer distinction between a model explicitly configured with zero experts and a model where expert configuration is absent.
• Accurate MoE Grouped GEMM Activation: Updated the condition for moe_grouped_gemm to activate only when num_experts is explicitly set (i.e., not None), preventing unintended MoE grouping for dense models that do not utilize experts.

🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console.